from .base_handler import BaseHandler
from tjob.initialize.setup_envs import (
    DEFAULT_DOCKER,
    _CACHE,
    _LOCAL_CONF,
    _DEBUG_MODE,
)
from tjob.utils import dump_exec_file, gen_uuid, get_curtime, load_yaml, JobInfo
import subprocess
import os.path as osp


class LocalHandler(BaseHandler):
    def __init__(self, name, worker_image) -> None:
        super().__init__()
        # use DEFAULT_DOCKER if image is None
        self._local_params = load_yaml(_LOCAL_CONF)
        self._image = (
            worker_image
            if worker_image
            else self._local_params.get("worker_image", DEFAULT_DOCKER)
        )
        self._name = name
        # gen id each times
        self._container_name = name if _DEBUG_MODE else gen_uuid()
        if _DEBUG_MODE:
            # keep docker
            self._cmd = (
                f"sudo docker run --name {self._container_name} "
                + "--gpus all -v /home:/home -v /mnt:/mnt --shm-size 2g "
            )
        else:
            self._cmd = (
                f"sudo docker run --name {self._container_name} --rm "
                + "--gpus all -v /home:/home -v /mnt:/mnt --shm-size 2g "
            )

        if self._local_params.get("interactive", True):
            self._cmd += "-it "

        self._cmd += self._image
        self._submit_info = JobInfo(
            name, "local", get_curtime(), id=self._container_name
        )

    def submit(self, command: str) -> int:
        # TODO(yuezhang): use jobname for consist?
        local_run_file = osp.join(_CACHE, "local_run.sh")
        dump_exec_file(command, local_run_file)
        # os.chmod(local_run_file, 755)
        self._cmd = self._cmd.strip().split(" ")
        if _DEBUG_MODE:
            print(local_run_file, "is not running with DEBUG mode")
        else:
            self._cmd.append(f"{local_run_file}")
        ret = subprocess.call(self._cmd)
        return ret

    @staticmethod
    def single_stop(id):
        stop_cmd = f"sudo docker stop {id}"
        stop_cmd = stop_cmd.split(" ")
        # need inputs passwd
        ret = subprocess.call(stop_cmd)
        return ret
